{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cee2ea73",
   "metadata": {},
   "source": [
    "# Training a Dense Neural Network with PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfe9838d",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820de326",
   "metadata": {},
   "source": [
    "This notebook was created by [Jupyter AI](https://github.com/jupyterlab/jupyter-ai) with the following prompt:\n",
    "\n",
    "> /generate A Jupyter notebook on training a dense neural network with 3 layers using PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a856dc73",
   "metadata": {},
   "source": [
    "This Jupyter notebook provides a detailed guide on training a dense neural network with 3 layers using PyTorch. The notebook covers the necessary steps, such as importing libraries, loading and preprocessing the dataset, defining the neural network architecture, initializing model parameters, defining the loss function and optimizer, training the neural network, and evaluating the trained model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "648a23da",
   "metadata": {},
   "source": [
    "## Loading and preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c19d6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9985df35",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f280be7",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2467f1b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80a40f98",
   "metadata": {},
   "source": [
    "## Defining the neural network architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26c0402",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defa866e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc3 = nn.Linear(hidden_size, output_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65bcea65",
   "metadata": {},
   "outputs": [],
   "source": [
    "    def forward(self, x):\n",
    "        out = self.fc1(x)\n",
    "        out = self.relu(out)\n",
    "        out = self.fc2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.fc3(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f29dff44",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_size = 100\n",
    "hidden_size = 50\n",
    "output_size = 10\n",
    "model = NeuralNetwork(input_size, hidden_size, output_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d93a016",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ce6136",
   "metadata": {},
   "source": [
    "## Initializing the model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "511e043e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1076cb99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the dimensions of the input and output layers\n",
    "input_size = 784  # Number of input features\n",
    "hidden_size = 128  # Number of neurons in the hidden layer\n",
    "output_size = 10  # Number of output classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70d16a55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the model parameters\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(input_size, hidden_size),  # Input layer to hidden layer\n",
    "    nn.ReLU(),  # Activation function\n",
    "    nn.Linear(hidden_size, output_size)  # Hidden layer to output layer\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8fafa4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Return the model architecture\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a144a416",
   "metadata": {},
   "source": [
    "## Defining the loss function and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd3e5a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb41639",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model\n",
    "model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b573826e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss function\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb385216",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86861f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the loss function and optimizer\n",
    "print(\"Loss Function:\", criterion)\n",
    "print(\"Optimizer:\", optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81a4929f",
   "metadata": {},
   "source": [
    "## Training the neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a2c514e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20bffe57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the training function\n",
    "def train(model, train_loader, criterion, optimizer, device):\n",
    "    # Set the model to train mode\n",
    "    model.train()\n",
    "    \n",
    "    # Loop through the training loader\n",
    "    for inputs, labels in train_loader:\n",
    "        # Move inputs and labels to the device\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        # Clear the gradients\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model(inputs)\n",
    "        \n",
    "        # Calculate the loss\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # Backward pass\n",
    "        loss.backward()\n",
    "        \n",
    "        # Update the weights\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7064b1d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the validation function\n",
    "def validate(model, val_loader, criterion, device):\n",
    "    # Set the model to eval mode\n",
    "    model.eval()\n",
    "    \n",
    "    # Initialize variables for tracking accuracy and loss\n",
    "    total_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    # Disable gradient calculation\n",
    "    with torch.no_grad():\n",
    "        # Loop through the validation loader\n",
    "        for inputs, labels in val_loader:\n",
    "            # Move inputs and labels to the device\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            # Forward pass\n",
    "            outputs = model(inputs)\n",
    "            \n",
    "            # Calculate the loss\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            # Update total loss\n",
    "            total_loss += loss.item()\n",
    "            \n",
    "            # Get the predicted labels\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            \n",
    "            # Update total and correct predictions\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    \n",
    "    # Calculate accuracy and average loss\n",
    "    accuracy = correct / total\n",
    "    avg_loss = total_loss / len(val_loader)\n",
    "    \n",
    "    return accuracy, avg_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41d4067c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set device (cuda if available, else cpu)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad96ff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set hyperparameters\n",
    "num_epochs = 10\n",
    "learning_rate = 0.001\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f32bd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create train and validation loaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280a1646",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model instance\n",
    "model = MyNeuralNetwork()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6303ff26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Move model to device\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85ad91ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f57fc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    # Train the model\n",
    "    train(model, train_loader, criterion, optimizer, device)\n",
    "    \n",
    "    # Validate the model\n",
    "    accuracy, avg_loss = validate(model, val_loader, criterion, device)\n",
    "    \n",
    "    # Print progress\n",
    "    print(f\"Epoch: {epoch+1}/{num_epochs} | Loss: {avg_loss:.4f} | Accuracy: {accuracy*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b15d75b",
   "metadata": {},
   "source": [
    "## Evaluating the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c3704b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39495a68",
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_predictions = 0\n",
    "total_samples = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "347129a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total_samples += labels.size(0)\n",
    "        correct_predictions += (predicted == labels).sum().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa4b20a",
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy = 100 * correct_predictions / total_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea30f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Accuracy: {accuracy:.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
